from tensorflow.keras.callbacks import Callback
import tensorflow as tf
import matplotlib.pyplot as plt
from IPython.display import clear_output
from PIL import Image
import numpy as np
import math

plt.figure(figsize=(20,10))


def display_preds(preds,epoch,totoal_epochs=10):
  rows = int(math.sqrt(totoal_epochs))
  cols = rows + 1  
  plt.subplot(rows,cols,epoch+1)
  plt.title(f'epoch: {epoch}')
  plt.imshow(tf.keras.preprocessing.image.array_to_img(preds))
  plt.axis('off')
  plt.tight_layout()
  plt.savefig(f'res/{epoch:04}.jpg')


def display(display_list,epoch):
  plt.figure(figsize=(15, 15))

  title = ['Input Image', 'True Mask', 'Predicted Mask']

  for i in range(len(display_list)):
    plt.subplot(1, len(display_list), i+1)
    plt.title(f'{epoch}: {title[i]}')
    plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))
    plt.axis('off')

  plt.savefig(f'res/{epoch:04}.jpg')


def create_mask(pred_mask):
  pred_mask = tf.argmax(pred_mask, axis=-1)
  pred_mask = pred_mask[..., tf.newaxis]
  return pred_mask[0]

def normalize(img):
    max_val = np.max(img)
    min_val = np.min(img)
    val_range = max_val - min_val
    norm_0_1 = (img-min_val)/val_range
    img = np.clip(2*norm_0_1-1,-1,1)
    return img

def show_predictions(model,epoch,target_size=[224,224]):
  image = Image.open("C:/Users/hblee/Documents/datasets/seg2d/val/JPEGImages/ISIC_0000027.jpg").convert('RGB')
  label = Image.open("C:/Users/hblee/Documents/datasets/seg2d/val/Segmentations/ISIC_0000027.png")
  # 操作图像
  image = image.resize(target_size,Image.BICUBIC)
  image_array = np.array(image).astype('float32')
  image_array = normalize(image_array) # 归一化
  image_array_p = np.expand_dims(image_array,axis=0)
  # 操作label
  label = label.resize(target_size,Image.NEAREST)
  label_array = np.expand_dims(np.array(label),axis=-1)

  pred_mask = model.predict(image_array_p)
  display_preds(create_mask(pred_mask),epoch)
  # display([image_array, label_array, create_mask(pred_mask)],epoch)  


class DisplayCallback(Callback):
  def __init__(self,model):
    self.model = model
  def on_epoch_end(self, epoch, logs=None):
    clear_output(wait=True)
    show_predictions(self.model,epoch)
    print ('\nepoch {}的预测结束, 保存在res文件夹中\n'.format(epoch+1))

class SaveModelCallback(Callback):
  def __init__(self,model):
      self.model = model
  
  def on_epoch_end(self, epoch, logs=None):
      self.model.save('saved_model/epoch')
      print(f'完成保存epoch: {epoch}的模型.')
